#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>

#define TILE_SIZE 16

// CUDA kernel for tiled matrix multiplication with boundary handling
__global__ void matrixMulTiled(float *A, float *B, float *C, int M, int K, int N)
{
    __shared__ float shared_A[TILE_SIZE][TILE_SIZE];
    __shared__ float shared_B[TILE_SIZE][TILE_SIZE];

    // Calculate thread row and column in the global matrix
    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;

    float value = 0.0;

    // Loop over tiles
    for (int tileIdx = 0; tileIdx < (K + TILE_SIZE - 1) / TILE_SIZE; ++tileIdx)
    {
        // Load elements of A and B into shared memory with boundary checks
        int a_idx = row * K + tileIdx * TILE_SIZE + threadIdx.x;
        int b_idx = (tileIdx * TILE_SIZE + threadIdx.y) * N + col;

        if (row < M && (tileIdx * TILE_SIZE + threadIdx.x) < K)
            shared_A[threadIdx.y][threadIdx.x] = A[a_idx];
        else
            shared_A[threadIdx.y][threadIdx.x] = 0.0f;

        if (col < N && (tileIdx * TILE_SIZE + threadIdx.y) < K)
            shared_B[threadIdx.y][threadIdx.x] = B[b_idx];
        else
            shared_B[threadIdx.y][threadIdx.x] = 0.0f;

        __syncthreads();

        // Perform the dot product for the tile
        for (int k = 0; k < TILE_SIZE; ++k)
        {
            value += shared_A[threadIdx.y][k] * shared_B[k][threadIdx.x];
        }

        __syncthreads();
    }

    // Store the result in C with boundary check
    if (row < M && col < N)
    {
        C[row * N + col] = value;
    }
}

int main()
{
    int M = 512; // Number of rows in A and C
    int K = 256; // Number of columns in A and number of rows in B
    int N = 512; // Number of columns in B and C

    size_t size_A = M * K * sizeof(float);
    size_t size_B = K * N * sizeof(float);
    size_t size_C = M * N * sizeof(float);

    float *h_A = (float *)malloc(size_A);
    float *h_B = (float *)malloc(size_B);
    float *h_C = (float *)malloc(size_C);

    for (int i = 0; i < M * K; ++i)
    {
        h_A[i] = rand() % 100 / 100.0f;
    }
    for (int i = 0; i < K * N; ++i)
    {
        h_B[i] = rand() % 100 / 100.0f;
    }

    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, size_A);
    cudaMalloc(&d_B, size_B);
    cudaMalloc(&d_C, size_C);

    cudaMemcpy(d_A, h_A, size_A, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, size_B, cudaMemcpyHostToDevice);

    dim3 threads(TILE_SIZE, TILE_SIZE);
    dim3 blocks((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);

    matrixMulTiled<<<blocks, threads>>>(d_A, d_B, d_C, M, K, N);

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess)
    {
        fprintf(stderr, "Failed to launch kernel: %s\n", cudaGetErrorString(err));
        exit(EXIT_FAILURE);
    }

    cudaMemcpy(h_C, d_C, size_C, cudaMemcpyDeviceToHost);

    err = cudaGetLastError();
    if (err != cudaSuccess)
    {
        fprintf(stderr, "Failed to copy memory from device to host: %s\n", cudaGetErrorString(err));
        exit(EXIT_FAILURE);
    }

    // Free memory
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    free(h_A);
    free(h_B);
    free(h_C);

    return 0;
}
